#使用此功能时,先调用 stocks_info_insert.py 导入股票信息,再运行脚本可一键更新所有股票日线数据。
import pandas as pd
from sqlalchemy import create_engine
import tushare as ts
from datetime import datetime, timedelta

# 获取电脑当日时间并转换为 YYYYMMDD 格式
end_date = datetime.now().strftime('%Y%m%d')

# 创建数据库连接引擎
engine = create_engine('mysql+pymysql://root:@127.0.0.1:3306/stocks?charset=utf8mb4&use_unicode=1')

# 从 stocks_info 表中读取所有 ts_code 和对应的 name
query_info = "SELECT ts_code, name FROM stocks_info"
info_df = pd.read_sql(query_info, con=engine)
info_ts_codes = info_df['ts_code'].tolist()
# 建立 ts_code 到 name 的映射
code_to_name = dict(zip(info_df['ts_code'], info_df['name']))

# 从 stocks_daily 表中读取所有 ts_code
query_daily = "SELECT DISTINCT ts_code FROM stocks_daily"
daily_ts_codes = pd.read_sql(query_daily, con=engine)['ts_code'].tolist()

# 找出不在 stocks_daily 表中的 ts_code
missing_ts_codes = [ts_code for ts_code in info_ts_codes if ts_code not in daily_ts_codes]

# 初始化 pro 接口
pro = ts.pro_api('')#已过期，请自行申请替换自己的token

# 处理不在 stocks_daily 表中的 ts_code
for insert_ts_codes in missing_ts_codes:
    start_date = '20240101'
    try:
        df = pro.daily(
            ts_code=insert_ts_codes,
            start_date=start_date,
            end_date=end_date
        )
    except Exception as e:
        print(f"拉取 {insert_ts_codes}({code_to_name.get(insert_ts_codes)}) 数据时出现错误: {e}")
    else:
        # 确保列名和数据库表结构一致
        db_columns = ['ts_code', 'trade_date', 'open', 'high', 'low', 'close',
                      'pre_close', 'change', 'pct_chg', 'vol', 'amount']

        # 只保留需要的列
        df = df[db_columns] if set(db_columns).issubset(df.columns) else df

        # 重命名 change 字段为 chg_amo 以匹配数据库
        if 'change' in df.columns:
            df = df.rename(columns={'change': 'chg_amo'})

        # 将 trade_date 转换为数据库需要的格式
        if 'trade_date' in df.columns:
            df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d').dt.date

        # 将新数据写入 stocks_daily 表
        if not df.empty:
            try:
                # 先检查是否有重复数据（双重保险）
                existing_dates_query = f"SELECT trade_date FROM stocks_daily WHERE ts_code = '{insert_ts_codes}'"
                existing_dates = pd.read_sql(existing_dates_query, con=engine)['trade_date'].tolist()
                df = df[~df['trade_date'].isin(existing_dates)]

                if not df.empty:
                    df.to_sql('stocks_daily', con=engine, if_exists='append', index=False)
                    print(
                        f"成功为 {insert_ts_codes}({code_to_name.get(insert_ts_codes)}) 写入 {len(df)} 条新数据，日期范围: {df['trade_date'].min()} 到 {df['trade_date'].max()}")
                else:
                    print(f"{insert_ts_codes}({code_to_name.get(insert_ts_codes)}) 没有需要更新的新数据")
            except Exception as e:
                print(f"为 {insert_ts_codes}({code_to_name.get(insert_ts_codes)}) 写入数据时出现错误: {e}")
        else:
            print(f"{insert_ts_codes}({code_to_name.get(insert_ts_codes)}) 在 {start_date} 到 {end_date} 期间没有新数据。")

# 遍历 stocks_daily 表中的每个 ts_code
for ts_code in daily_ts_codes:
    # 查询该 ts_code 的最新 trade_date
    query = f"SELECT MAX(trade_date) as latest_date FROM stocks_daily WHERE ts_code = '{ts_code}'"
    latest_date_df = pd.read_sql(query, con=engine)
    latest_date = latest_date_df['latest_date'].values[0]
    print(f"最新日期: {latest_date}")
    # 处理 latest_date 为 YYYYMMDD 格式的字符串
    start_date = (latest_date + timedelta(days=1)).strftime('%Y%m%d')

    # 检查 start_date 是否大于 end_date（避免无效查询）
    if start_date > end_date:
        print(f"{ts_code}({code_to_name.get(ts_code)}) 已是最新数据，无需更新")
        continue

    # 拉取新数据
    try:
        df = pro.daily(
            ts_code=ts_code,
            start_date=start_date,
            end_date=end_date
        )
    except Exception as e:
        print(f"拉取 {ts_code}({code_to_name.get(ts_code)}) 数据时出现错误: {e}")
        continue

    # 确保列名和数据库表结构一致
    db_columns = ['ts_code', 'trade_date', 'open', 'high', 'low', 'close',
                  'pre_close', 'change', 'pct_chg', 'vol', 'amount']

    # 只保留需要的列
    df = df[db_columns] if set(db_columns).issubset(df.columns) else df

    # 重命名 change 字段为 chg_amo 以匹配数据库
    if 'change' in df.columns:
        df = df.rename(columns={'change': 'chg_amo'})

    # 将 trade_date 转换为数据库需要的格式
    if 'trade_date' in df.columns:
        df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d').dt.date

    # 将新数据写入 stocks_daily 表
    if not df.empty:
        try:
            # 先检查是否有重复数据（双重保险）
            existing_dates_query = f"SELECT trade_date FROM stocks_daily WHERE ts_code = '{ts_code}'"
            existing_dates = pd.read_sql(existing_dates_query, con=engine)['trade_date'].tolist()
            df = df[~df['trade_date'].isin(existing_dates)]

            if not df.empty:
                df.to_sql('stocks_daily', con=engine, if_exists='append', index=False)
                print(
                    f"成功为 {ts_code}({code_to_name.get(ts_code)}) 写入 {len(df)} 条新数据，日期范围: {df['trade_date'].min()} 到 {df['trade_date'].max()}")
            else:
                print(f"{ts_code}({code_to_name.get(ts_code)}) 没有需要更新的新数据")
        except Exception as e:
            print(f"为 {ts_code}({code_to_name.get(ts_code)}) 写入数据时出现错误: {e}")
    else:
        print(f"{ts_code}({code_to_name.get(ts_code)}) 在 {start_date} 到 {end_date} 期间没有新数据。")

